from Mapillary.mapillary_sls.utils.utils import configure_transform
from Mapillary.mapillary_sls.datasets.msls import MSLS

def loadmapillary(datapath, mode="train"):
    meta = {'mean': [0.485, 0.456, 0.406], 'std': [0.229, 0.224, 0.225]}
    transform = configure_transform(image_dim = (480, 640), meta = meta)

    posDistThr = 5

    # negatives are defined outside a radius of 25 m
    negDistThr = 25

    # number of negatives per triplet
    nNeg = 20

    # number of cached queries
    cached_queries = 1

    # number of cached negatives
    cached_negatives = 100

    # whether to use positive sampling
    positive_sampling = True

    # choose the cities to load
    cities = 'zurich'

    # choose task to test on [im2im, seq2im, im2seq, seq2seq]
    task = 'im2im'

    # choose sequence length
    seq_length = 1

    # my_root_dir = '/workspace/wzj/dataset/image_retrieval/Mapillary'
    if mode == 'train':
        dataset = MSLS(datapath, cities = cities, transform = transform, mode = "train", task = task, seq_length = seq_length,
                        negDistThr = negDistThr, posDistThr = posDistThr, nNeg = nNeg, cached_queries = cached_queries,
                        cached_negatives = cached_negatives, positive_sampling = positive_sampling)
    elif mode == 'val':
        dataset = MSLS(datapath, cities = cities, transform = transform, mode = 'val',
                   task = task, seq_length = seq_length, subtask = 'all', posDistThr = posDistThr)

    return dataset